# -*- coding: UTF-8 -*-
# *******************************************************************
# File Name: read_mtwi
# > Author: 04000387
# > Created Time: 2024/12/23 16:34
# *******************************************************************
import cv2
from pathlib import Path
from torch.utils.data import Dataset


class ReadMtwi(Dataset):
    def __init__(self, save_dir):
        self.path = Path(save_dir)
        self.img_path = self.path / "image_train"
        self.text_path = self.path / "txt_train"
        # 检查文件夹是否存在
        self._check_exist()

        # 获取img
        self.img_files = self.get_files()

    def _check_exist(self):
        if not (self.path.exists() and self.img_path.exists() and self.text_path.exists()):
            msg = self.path.absolute()
            raise ValueError("文件夹:[{}]不存在".format(msg))

    def get_files(self):
        res = []
        for it in self.img_path.iterdir():
            res.append(it)

        return res

    def deal_text(self, path):
        with open(path, "r", encoding="utf-8") as fp:
            lines = fp.readlines()

        boxes = []
        text = []
        for line in lines:
            line = line.strip().strip("\n")
            items = line.split(",")
            box_one = [(float(items[0]), float(items[1])),
                       (float(items[2]), float(items[3])),
                       (float(items[4]), float(items[5])),
                       (float(items[6]), float(items[7]))
                       ]

            text_one = "".join(items[8:])

            boxes.append(box_one)
            text.append(text_one)
        return {"boxes": boxes, "text": text}

    def __len__(self):
        return len(self.img_files)

    def __getitem__(self, idx):
        img_path = self.img_files[idx]
        img = cv2.imread(str(img_path.absolute()))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        text_path = self.text_path / (img_path.stem + ".txt")
        box_info = self.deal_text(text_path)
        return img, box_info
